#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# 文档关键字搜索工具 - 全盘搜索版（兼容 Python 3.5）

import os
import re
import tkinter as tk
from tkinter import ttk, messagebox
import threading
import subprocess
import sys

# ==================== 配置区域 ====================

# 搜索根目录：全盘搜索
SEARCH_ROOTS = [
    "/home",
    "/root",
    "/opt",
    "/tmp",
    "/usr/local",
    "/etc"
]

# 支持的文件类型
SUPPORTED_EXTENSIONS = {
    '.txt': 'text',
    '.doc': 'doc',
    '.docx': 'docx',
    '.xls': 'xls',
    '.xlsx': 'xlsx',
    '.ppt': 'ppt',
    '.pptx': 'pptx',
    '.pdf': 'pdf'
}

# 最大年份限制（用于检测新文件）
MAX_YEAR = 2025

# 分隔符
SEPARATOR = "-" * 60

# 系统/权限受限目录（跳过）
SKIP_DIRS = {
    '/proc', '/sys', '/dev', '/run', '/tmp', '/var/run',
    '/boot', '/lost+found', '/mnt', '/media', '/snap',
    '/usr/share', '/lib', '/lib64', '/usr/lib', '/usr/lib64'
}

# 默认关键字
DEFAULT_KEYWORDS = "秘密,机密,绝密"

# ==================== 工具函数 ====================

def safe_tcl_string(s):
    """移除 Tcl 不支持的高位 Unicode 字符（U+10000 以上）"""
    s = str(s)
    return ''.join(c for c in s if ord(c) <= 0xFFFF)

# ==================== 提取器函数 ====================

def extract_text(filepath):
    try:
        with open(filepath, 'r', encoding='utf-8', errors='ignore') as f:
            return f.read()
    except:
        return ""

def extract_doc(filepath):
    try:
        # 使用 antiword 提取 .doc
        result = subprocess.run(['antiword', filepath], capture_output=True, text=True, timeout=10)
        return result.stdout
    except Exception as e:
        return ""

def extract_docx(filepath):
    try:
        import docx
        doc = docx.Document(filepath)
        return "\n".join(paragraph.text for paragraph in doc.paragraphs)
    except:
        return ""

def extract_xls(filepath):
    try:
        import xlrd
        workbook = xlrd.open_workbook(filepath)
        text = ""
        for sheet in workbook.sheets():
            for row in range(sheet.nrows):
                for col in range(sheet.ncols):
                    cell_value = sheet.cell_value(row, col)
                    text += str(cell_value) + "\t"
                text += "\n"
        return text
    except:
        return ""

def extract_xlsx(filepath):
    try:
        import openpyxl
        workbook = openpyxl.load_workbook(filepath, read_only=True, data_only=True)
        text = ""
        for sheet in workbook.worksheets:
            for row in sheet.rows:
                text += "\t".join([str(cell.value) if cell.value is not None else "" for cell in row]) + "\n"
        return text
    except:
        return ""

def extract_ppt(filepath):
    try:
        # 尝试使用 strings 提取文本（适用于旧版 .ppt）
        result = subprocess.run(['strings', filepath], capture_output=True, text=True, timeout=10)
        return result.stdout
    except Exception as e:
        return ""

def extract_pptx(filepath):
    try:
        from pptx import Presentation
        pres = Presentation(filepath)
        text = ""
        for slide in pres.slides:
            for shape in slide.shapes:
                if hasattr(shape, "text"):
                    text += shape.text + "\n"
        return text
    except:
        return ""

def extract_pdf(filepath):
    try:
        result = subprocess.run(['pdftotext', '-q', filepath, '-'], capture_output=True, text=True, timeout=10)
        return result.stdout
    except:
        return ""

# 所有提取器映射
EXTRACTORS = {
    'text': extract_text,
    'doc': extract_doc,
    'docx': extract_docx,
    'xls': extract_xls,
    'xlsx': extract_xlsx,
    'ppt': extract_ppt,
    'pptx': extract_pptx,
    'pdf': extract_pdf,
}

# 获取文件创建年份
def get_creation_year(filepath):
    try:
        return int(os.path.getctime(filepath)) // (365*24*3600) + 1970
    except:
        return None

# ==================== GUI 主程序 ====================

class DocumentSearcher:
    def __init__(self, root):
        self.root = root
        self.root.title("文档关键字搜索工具 - 全盘搜索版")
        self.root.geometry("900x700")

        self.stop_searching = threading.Event()
        self.search_thread = None

        self.setup_ui()

        # 设置默认关键字
        self.keyword_entry.insert(0, DEFAULT_KEYWORDS)

    def setup_ui(self):
        frame = ttk.Frame(self.root, padding="10")
        frame.grid(row=0, column=0, sticky=(tk.W, tk.E, tk.N, tk.S))

        # 关键字输入
        ttk.Label(frame, text="关键字（逗号分隔）：").grid(row=0, column=0, sticky=tk.W)
        self.keyword_entry = ttk.Entry(frame, width=50)
        self.keyword_entry.grid(row=0, column=1, sticky=tk.W, padx=(5, 0))

        # 按钮
        self.search_button = ttk.Button(frame, text="开始搜索", command=self.start_search)
        self.search_button.grid(row=0, column=2, padx=(10, 0))

        self.stop_button = ttk.Button(frame, text="停止搜索", command=self.stop_search, state=tk.DISABLED)
        self.stop_button.grid(row=0, column=3, padx=(5, 0))

        # 日志显示（在上）
        log_frame = ttk.Frame(self.root)
        log_frame.grid(row=1, column=0, sticky=(tk.W, tk.E, tk.N, tk.S), padx=10, pady=(5, 2))
        ttk.Label(log_frame, text="日志信息：").pack(anchor=tk.W)

        self.log_text = tk.Text(log_frame, height=12, wrap=tk.WORD, fg="gray", font=("Consolas", 9))
        self.log_text.pack(side=tk.LEFT, fill=tk.BOTH, expand=True)

        log_scroll = ttk.Scrollbar(log_frame, orient=tk.VERTICAL, command=self.log_text.yview)
        log_scroll.pack(side=tk.RIGHT, fill=tk.Y)
        self.log_text.config(yscrollcommand=log_scroll.set)

        # 结果显示（在下）
        result_frame = ttk.Frame(self.root)
        result_frame.grid(row=2, column=0, sticky=(tk.W, tk.E, tk.N, tk.S), padx=10, pady=(2, 10))
        ttk.Label(result_frame, text="搜索结果：").pack(anchor=tk.W)

        self.result_text = tk.Text(result_frame, height=20, wrap=tk.WORD, font=("Consolas", 10))
        self.result_text.pack(side=tk.LEFT, fill=tk.BOTH, expand=True)

        result_scroll = ttk.Scrollbar(result_frame, orient=tk.VERTICAL, command=self.result_text.yview)
        result_scroll.pack(side=tk.RIGHT, fill=tk.Y)
        self.result_text.config(yscrollcommand=result_scroll.set)

        # ==================== 新增：底部版权标签 ====================
        footer_frame = ttk.Frame(self.root)
        footer_frame.grid(row=3, column=0, sticky=tk.W+tk.E, padx=10, pady=(5, 10))

        # 居中容器
        footer_label = ttk.Label(
            footer_frame,
            text="文山州公安局刑侦支队 张泽山@2025",
        )
        footer_label.pack(side=tk.TOP, anchor=tk.CENTER, expand=True)

        # 强制居中（通过 grid columnconfigure）
        footer_frame.columnconfigure(0, weight=1)
        # ==============================================================

        # 配置权重
        self.root.columnconfigure(0, weight=1)
        self.root.rowconfigure(1, weight=1)
        self.root.rowconfigure(2, weight=3)
        # 注意：row 3（footer）不设置 weight，防止拉伸

    def log_message(self, msg):
        msg = safe_tcl_string(msg)
        self.log_text.insert(tk.END, msg + "\n")
        self.log_text.see(tk.END)
        self.log_text.update_idletasks()

    def start_search(self):
        keywords = self.keyword_entry.get().strip()
        if not keywords:
            messagebox.showwarning("警告", "请输入关键字！")
            return

        self.result_text.delete(1.0, tk.END)
        self.search_button.config(state=tk.DISABLED)
        self.stop_button.config(state=tk.NORMAL)
        self.stop_searching.clear()

        keywords_list = [k.strip() for k in keywords.split(',') if k.strip()]
        self.log_message("🔍 开始搜索关键字: {}".format(', '.join(keywords_list)))
        self.log_message("📁 搜索路径: {}".format(', '.join(SEARCH_ROOTS)))

        def worker():
            found_files = set()

            for root_dir in SEARCH_ROOTS:
                if self.stop_searching.is_set():
                    break
                self.log_message("➡️ 开始遍历目录: {}".format(root_dir))

                try:
                    for dirpath, dirnames, filenames in os.walk(root_dir, followlinks=True):
                        if self.stop_searching.is_set():
                            break

                        # 跳过系统/权限受限目录
                        if any(dirpath.startswith(skip) for skip in SKIP_DIRS):
                            continue

                        try:
                            self.log_message("📁 进入目录: {} (文件数: {})".format(dirpath, len(filenames)))
                        except:
                            pass  # 防止路径编码问题

                        for filename in filenames:
                            if self.stop_searching.is_set():
                                break

                            filepath = os.path.join(dirpath, filename)
                            if not os.path.isfile(filepath):
                                continue

                            # 检查年份（可选）
                            try:
                                year = get_creation_year(filepath)
                                if year is not None and year > MAX_YEAR:
                                    self.log_message("💡 本软件版本已停用，请下载新版本使用")
                                    self.stop_searching.set()
                                    return
                            except:
                                pass

                            _, ext = os.path.splitext(filename)
                            ext = ext.lower()
                            if ext not in SUPPORTED_EXTENSIONS:
                                continue

                            file_type = SUPPORTED_EXTENSIONS[ext]
                            extractor = EXTRACTORS.get(file_type)
                            if not extractor:
                                continue

                            try:
                                text = extractor(filepath)
                                if not text:
                                    continue

                                text_lower = text.lower()
                                matched_keywords = []
                                for kw in keywords_list:
                                    kw_lower = kw.lower()
                                    if kw_lower in text_lower:
                                        for match in re.finditer(re.escape(kw_lower), text_lower):
                                            start = max(0, match.start() - 50)
                                            end = min(len(text), match.end() + 50)
                                            context = text[start:end]
                                            matched_keywords.append((kw, context))

                                if matched_keywords:
                                    if filepath not in found_files:
                                        self.result_text.insert(tk.END, "{}\n{}\n".format(SEPARATOR, filepath))
                                        found_files.add(filepath)
                                    for mkw, ctx in matched_keywords:
                                        output = "  🔑 关键字: '{}'\n  📄 上下文: {}\n".format(mkw, ctx)
                                        self.result_text.insert(tk.END, safe_tcl_string(output))
                                    self.result_text.see(tk.END)
                                    self.result_text.update_idletasks()
                            except Exception as e:
                                self.log_message("❌ 处理文件失败: {} - {}".format(filepath, str(e)))

                except PermissionError:
                    self.log_message("⛔ 无权限访问目录: {}".format(root_dir))
                except Exception as e:
                    self.log_message("❌ 遍历目录失败: {} - {}".format(root_dir, str(e)))

            self.log_message("✅ 搜索完成！共找到 {} 个匹配文件。".format(len(found_files)))
            self.search_button.config(state=tk.NORMAL)
            self.stop_button.config(state=tk.DISABLED)
            self.stop_searching.clear()

        self.search_thread = threading.Thread(target=worker, daemon=True)
        self.search_thread.start()

    def stop_search(self):
        self.stop_searching.set()
        self.log_message("🛑 正在停止搜索...")
        self.stop_button.config(state=tk.DISABLED)

# ==================== 主程序入口 ====================

if __name__ == "__main__":
    root = tk.Tk()
    app = DocumentSearcher(root)
    root.mainloop()
